import torch
import os
import pickle
import numpy as np
def save_checkpoint(state, filename="checkpoint.pth.tar"):
    """保存当前的模型和训练状态"""
    torch.save(state, filename)

def load_checkpoint(checkpoint_path, model, optimizer):
    """从检查点文件加载模型和优化器状态，并返回历史数据和开始的epoch"""
    if os.path.isfile(checkpoint_path):
        print("=> Loading checkpoint")
        checkpoint = torch.load(checkpoint_path)
        model.load_state_dict(checkpoint['state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        history = checkpoint.get('history', {'train_loss': [], 'test_loss': []})  # 加载或初始化历史数据
        return checkpoint['epoch'], history
    else:
        print("=> No checkpoint found at '{}'".format(checkpoint_path))
        return 0, {'train_loss': [], 'test_loss': []}


def load_dataset_splits(split_datasets_file_idx):
    with open(split_datasets_file_idx, 'rb') as file:
        data = pickle.load(file)
    return data['train_imgs'], data['valid_imgs'], data['train_masks'], data['valid_masks']


def load_test_set(split_datasets_file_idx):
    with open(split_datasets_file_idx, 'rb') as file:
        data = pickle.load(file)
    return data['test_imgs'], data['test_masks']


def is_directory_empty(directory):
    """Check if a directory is empty."""
    if os.path.exists(directory):
        # 列出目录中的所有文件和子目录
        if not os.listdir(directory):  # 如果列表为空，返回True
            return True
        else:
            return False
    else:
        return True


def onehot_to_single_channel(onehot_array):

    # 使用argmax函数沿着第一个维度（通道维度）找到最大值的索引
    single_channel = np.argmax(onehot_array, axis=0)

    return single_channel